"""
split.py
功能：把 data/raw/ 下的 5 类图片按 8:2 随机分成训练集和测试集
运行：python split.py
"""
import os, shutil, random
from pathlib import Path
from src.config import Config  # 读取统一配置

def split():
    # 固定随机种子，保证每次划分结果相同
    random.seed(Config.seed)

    # 逐个类别处理
    for cls in os.listdir(Config.raw_dir):
        # 当前类别的完整路径
        src_dir = Path(Config.raw_dir) / cls

        # 获取该类别下所有图片文件（支持常见后缀）
        all_imgs = [f for f in src_dir.iterdir()
                    if f.suffix.lower() in ['.jpg', '.png', '.jpeg']]

        # 打乱顺序
        random.shuffle(all_imgs)

        # 计算训练集切分下标
        split_idx = int(len(all_imgs) * Config.split_ratio)

        # 分别处理训练集和测试集
        for mode, imgs in [('train', all_imgs[:split_idx]),
                           ('test',  all_imgs[split_idx:])]:
            # 目标目录：data/train/类别 或 data/test/类别
            dst_dir = Path('data') / mode / cls
            dst_dir.mkdir(parents=True, exist_ok=True)  # 若目录不存在则创建

            # 复制文件
            for img in imgs:
                shutil.copy(img, dst_dir / img.name)

    print('✅数据集划分完成！')

if __name__ == '__main__':
    split()